"""
六足机器人平衡控制器模块

该模块实现了六足机器人的平衡控制功能，包括：
- IMU数据采集和处理
- 姿态计算和补偿
- 平衡控制算法
"""

import math
import time
from typing import Dict, Any, Optional
import utils.filter as filter
from robot_core.config.servo_config import servo_config
from robot_core.control import MotionController
from sensors.imu_sensor import IMUSensor


class BalanceController:
    """六足机器人平衡控制器类"""

    # 运动学参数
    L = 220.00  # 前后腿距离
    W = 158.50  # 左右腿距离
    MW = 266.90  # 中间腿距离

    def __init__(self):
        """初始化平衡控制器"""
        self.imu = IMUSensor(0x79)
        self.imu.set_gyro_range(self.imu.GYRO_RANGE_2000DEG)
        self.imu.set_accel_range(self.imu.ACCEL_RANGE_2G)

        self.motion_controller = MotionController(servo_config)

        self.offset: Optional[Dict[str, float]] = None
        self.offset_row = 0.0
        self.offset_pit = 0.0
        self.offset_m = 0.0
        self.running = True

        # 初始化姿态
        self._init_balance()

    def _init_balance(self) -> None:
        """初始化平衡控制器"""
        # 初始化摄像头位置
        self.motion_controller._init_camera()

        # 设置坐下姿态
        self.motion_controller.sit()
        time.sleep(1)

        # 获取IMU偏差
        self.offset = self._get_offset()

        # 设置站立姿态
        self.motion_controller.stand()
        time.sleep(1)

    def _get_offset(self) -> Dict[str, float]:
        """
        获取IMU传感器的初始偏差

        Returns:
            包含加速度和角速度偏差的字典
        """
        offset_ax = 0.0
        offset_ay = 0.0
        offset_az = 0.0
        offset_gx = 0.0
        offset_gy = 0.0
        offset_gz = 0.0

        # 测量100次取平均值
        for _ in range(100):
            accel_data = self.imu.get_accel_data(g=True)
            gyro_data = self.imu.get_gyro_data()

            offset_ax += accel_data["x"]
            offset_ay += accel_data["y"]
            offset_az += accel_data["z"]
            offset_gx += gyro_data["x"]
            offset_gy += gyro_data["y"]
            offset_gz += gyro_data["z"]

        return {
            "ax": offset_ax / 100,
            "ay": offset_ay / 100,
            "az": offset_az / 100 - 1,
            "gx": offset_gx / 100,
            "gy": offset_gy / 100,
            "gz": offset_gz / 100,
        }

    def balance(self, row: float, pit: float, speed: int) -> None:
        """
        执行平衡控制

        Args:
            row: 横滚角(度)
            pit: 俯仰角(度)
            speed: 运动速度(ms)
        """
        if self.offset is None:
            raise RuntimeError("平衡控制器未初始化")

        # 避免小范围抖动
        if abs(row) < 0.5:
            row = 0
        if abs(pit) < 0.5:
            pit = 0

        # 角度转弧度
        row_rad = row * math.pi / 180
        pit_rad = pit * math.pi / 180

        # 计算各腿位置偏移
        self.offset_m += self.MW * math.tan(row_rad)
        self.offset_row += self.W * math.tan(row_rad)
        self.offset_pit += self.L * math.tan(pit_rad)

        # 限位检查
        if (
            abs(self.offset_m) > 100
            or abs(self.offset_row + self.offset_pit) > 100
            or abs(self.offset_pit - self.offset_row) > 100
        ):
            print("平衡控制超出限位")
            return

        # 计算各腿位置
        leg_positions = []
        for leg in range(6):
            position = [100.0, 100.0, -70.0]

            if leg == 0:
                position[2] += self.offset_row + self.offset_pit
            elif leg == 1:
                position[2] += self.offset_m
            elif leg == 2:
                position[2] += -self.offset_row + self.offset_pit
            elif leg == 3:
                position[2] += -self.offset_row - self.offset_pit
            elif leg == 4:
                position[2] += -self.offset_m
            elif leg == 5:
                position[2] += self.offset_row - self.offset_pit

            leg_positions.append(position)

        # 控制各腿运动
        for leg, position in enumerate(leg_positions):
            self.motion_controller.move_leg(leg, position, speed)

        time.sleep(speed * 0.001)

    def rolling(self) -> None:
        """执行滚动动作"""
        # 初始姿态
        self.balance(0.2094, 0.2094, 1000)
        time.sleep(1)

        # 执行圆周运动
        R = 0.2443
        while True:
            for i in range(20):
                A = i * 0.314
                row = R * math.cos(A)
                pit = R * math.sin(A)
                self.balance(row, pit, 100)
                time.sleep(0.1)

    def run(self) -> None:
        """运行平衡控制循环"""
        while True:
            if self.running:
                # 获取IMU数据
                accel_data = self.imu.get_accel_data(g=True)
                gyro_data = self.imu.get_gyro_data()

                # 计算姿态角
                ax = accel_data["x"] - self.offset["ax"]
                ay = accel_data["y"] - self.offset["ay"]
                az = accel_data["z"] - self.offset["az"]
                gx = gyro_data["x"] - self.offset["gx"]
                gy = gyro_data["y"] - self.offset["gy"]

                # 计算欧拉角
                angle_x = math.atan2(ax, az)
                angle_y = math.atan2(ay, az)

                # 使用互补滤波计算姿态
                angle_r = float("%.4f" % (filter.filter_r(angle_x * 180 / math.pi, gy)))
                angle_p = float("%.4f" % (filter.filter_p(angle_y * 180 / math.pi, gx)))

                # 执行平衡控制
                self.balance(angle_r, angle_p, 50)
                time.sleep(0.05)
            else:
                time.sleep(0.01)

    def stop(self) -> None:
        """停止平衡控制"""
        self.running = False

    def close(self) -> None:
        """关闭平衡控制器"""
        self.stop()


# 创建全局平衡控制器实例
balance_controller = BalanceController()
